[Feat]mtp aclgraph support #3244
Conversation
|
👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:
If CI fails, you can run linting and testing checks locally according Contributing and Testing. |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This PR adds ACLGraph support for the MTP model. The changes are correct in principle, but they expose a pre-existing critical bug in the speculative decoding loop in mtp_proposer.py. The batch_descriptor and other graph-related parameters are not updated across loop iterations, which will lead to incorrect model outputs when using more than one speculative token. I've left a comment with details on the issue and a suggested fix.
anon189Ty
force-pushed
the
mtp_aclgraph_support
branch
from
68201ac to
fd1994e
Compare
| self.drafter.dummy_run( | ||
| num_tokens=num_tokens, | ||
| with_prefill=with_prefill, | ||
| skip_attn=True, | ||
| num_reqs=num_reqs, | ||
| num_tokens_across_dp=num_tokens_across_dp) | ||
| num_tokens_across_dp=num_tokens_across_dp, | ||
| aclgraph_runtime_mode=aclgraph_runtime_mode, | ||
| batch_descriptor=batch_descriptor) |
There was a problem hiding this comment.
Remember to change the signature of dummy_run in class Proposer and all of its sub-classes
Signed-off-by: anon189Ty <[email protected]>
anon189Ty
force-pushed
the
mtp_aclgraph_support
branch
from
fd1994e to
58959a1
Compare
What this PR does / why we need it?
Currently, MTP Model in deepseek can not be capture in ACLGraph. This PR is use to allow MTP to be captured in ACLGraph mode.
Does this PR introduce any user-facing change?
How was this patch tested?